package com.mj.listen1._07AVL树.tree;

import java.util.Comparator;

/**
 * AVL 自平衡
 *
 * @param <E>
 */
public class AVLTree<E> extends BinarySearchTree<E> {
    public AVLTree() {

    }

    public AVLTree(Comparator<E> comparator) {
        super(comparator);
    }

    // 删除后的调整-自平衡调整
    @Override
    protected void afterRemove(Node<E> node) {
        // 找出失衡的节点——往父亲节点找
        while ((node = node.parent) != null) {
            // 此时的节点已经是父节点
            if (isBalance(node)) {
                // 是平衡的。更新高度
                updateHeight(node);
            } else {
                // 不是平衡的
                // 调整平衡
                reBalance(node);
            }
        }
    }

    @Override
    protected void afterAdd(Node<E> node) {
        // 找出失衡的节点——往父亲节点找
        while ((node = node.parent) != null) {
            // 此时的节点已经是父节点
            if (isBalance(node)) {
                // 是平衡的。更新高度
                updateHeight(node);
            } else {
                // 不是平衡的
                // 调整平衡
                reBalance(node);
                // 找到了失衡节点就不再去找了
                break;
            }
        }
    }

    /**
     * 统一调整平衡
     *
     * @param grand 最低的那个失去平衡的节点
     */
    private void reBalance(Node<E> grand) {
        // 父节点
        Node<E> parent = ((AVLNode<E>) grand).tallerChild();
        Node<E> node = ((AVLNode<E>) parent).tallerChild();
        // 判断方向
        if (parent.isLeftChild()) {// L
            if (node.isLeftChild()) {// LL
                rotate(grand, node.left, node, node.right, parent, parent.right, grand, grand.right);
            } else {// LR
                rotate(grand, parent.left, parent, node.left, node, node.right, grand, grand.right);
            }
        } else {// R
            if (node.isLeftChild()) {// RL
                rotate(grand, grand.left, grand, node.left, node, node.right, parent, parent.right);
            } else {// RR
                rotate(grand, grand.left, grand, parent.left, parent, node.left, node, node.right);
            }
        }
    }

    private void rotate(
            Node<E> root,// 原来子树的根节点
            Node<E> a, Node<E> b, Node<E> c,
            Node<E> d,
            Node<E> e, Node<E> f, Node<E> g) {
        // 让d成为这颗子树的根节点
        d.parent = root.parent;
        if (root.isLeftChild()) {
            root.parent.left = d;
        } else if (root.isRightChild()) {
            root.parent.right = d;
        } else {
            // 是根节点
            root = d;
        }

        // a-b-c
        b.left = a;
        if (a != null) {
            a.parent = b;
        }
        b.right = c;
        if (c != null) {
            c.parent = b;
        }
        updateHeight(b);

        // e-f-g
        f.left = e;
        if (e != null) {
            e.parent = f;
        }
        f.right = g;
        if (g != null) {
            g.parent = f;
        }
        updateHeight(f);

        // b-d-f
        d.left = b;
        d.right = f;
        b.parent = d;
        f.parent = d;
        updateHeight(d);
    }

    /**
     * 调整平衡
     *
     * @param grand 最低的那个失去平衡的节点
     */
    private void reBalance2(Node<E> grand) {
        // 父节点
        Node<E> parent = ((AVLNode<E>) grand).tallerChild();
        Node<E> node = ((AVLNode<E>) parent).tallerChild();
        // 判断方向
        if (parent.isLeftChild()) {// L
            if (node.isLeftChild()) {// LL
                rotateR(grand);
            } else {// LR
                rotateL(parent);
                rotateR(grand);
            }
        } else {// R
            if (node.isLeftChild()) {// RL
                rotateR(parent);
                rotateL(grand);
            } else {// RR
                rotateL(grand);
            }
        }
    }

    // 左旋
    private void rotateL(Node<E> grand) {
        Node<E> parent = grand.right;
        Node<E> child = parent.left;
        grand.right = child;
        parent.left = grand;
        afterRotate((Node<E>) grand, (Node<E>) parent, (Node<E>) child);

    }

    private void afterRotate(Node<E> grand, Node<E> parent, Node<E> child) {
        // 让 parent 为子树的根节点
        parent.parent = grand.parent;
        // 将grand.parent的左或者右指针指向parent
        // 如果 grand 是其父节点的 left, 则将 grand.parent.left = parent;
        if (grand.isLeftChild()) {
            grand.parent.left = parent;
        } else if (grand.isRightChild()) {
            grand.parent.right = parent;
        } else {
            // 根节点
            root = parent;
        }

        // 更新 child 的 parent
        if (child != null) {
            child.parent = grand;
        }

        // 更新grand的parent
        grand.parent = parent;

        // 更新高度(因为grand、parent的左右子树都变了, 先更新grand(矮的))
        updateHeight(grand);
        updateHeight(parent);
    }

    // 右旋
    private void rotateR(Node<E> grand) {
        Node<E> parent = grand.left;
        Node<E> child = parent.right;
        grand.left = child;
        parent.right = grand;

        // 让 parent 为子树的根节点
        afterRotate(grand, parent, child);
    }

    private void updateHeight(Node<E> node) {
        ((AVLNode<E>) node).updateHeight();
    }

    // AVL 生成自己的AVLNode
    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new AVLNode<>(element, parent);
    }

    // 判断当前节点是否是平衡的
    private boolean isBalance(Node<E> node) {
        return Math.abs(((AVLNode<E>) node).balanceFactor()) <= 1;
    }

    protected static class AVLNode<E> extends Node<E> {
        // 该节点的高度
        int height = 1;

        public AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }

        // 获取平衡因子
        public int balanceFactor() {
            // 左孩子的高度
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;

            // 右孩子的高度
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;

            // 差值
            return leftHeight - rightHeight;
        }

        public void updateHeight() {
            // 左孩子的高度
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            // 右孩子的高度
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            // 更新节点高度
            height = 1 + Math.max(leftHeight, rightHeight);
        }

        public Node<E> tallerChild() {
            // 左孩子的高度
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            // 右孩子的高度
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            // 返回孩子中高度最高的节点
            if (leftHeight < rightHeight) return right;
            if (leftHeight > rightHeight) return left;
            // 相等。自己是父节点的哪边就返回哪边的孩子
            return isLeftChild() ? left : right;
        }
    }

    @Override
    public Object string(Object node) {
        AVLNode<E> myNode = (AVLNode<E>) node;
        String parentStr = "null";
        if (myNode.parent != null) {
            parentStr = myNode.parent.element.toString();
        }
        return myNode.element + "_p(" + parentStr + ")" + "_h(" + myNode.height + ")";
    }
}
